Diagonal traverse II¶
Time: O(MxN); Space: O(M); medium
Given a list of lists of integers, nums, return all elements of nums in diagonal order as shown in the below images.
Example 1:
Input: nums = [[1,2,3],[4,5,6],[7,8,9]]
Output: [1,4,2,7,5,3,8,6,9]
Example 2:
Input: nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
Output: [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]
Example 3:
Input: nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]]
Output: [1,4,2,5,3,8,6,9,7,10,11]
Example 4:
Input: nums = [[1,2,3,4,5,6]]
Output: [1,2,3,4,5,6]
Constraints:
1 <= len(nums) <= 10^5
1 <= len(nums[i]) <= 10^5
1 <= nums[i][j] <= 10^9
There at most 10^5 elements in nums.
Hints:
Notice that numbers with equal sums of row and column indexes belong to the same diagonal.
Store them in tuples (sum, row, val), sort them, and then regroup the answer.
[1]:
import collections
class Solution1(object):
"""
Time: O(M*N)
Space: O(M)
"""
def findDiagonalOrder(self, nums):
"""
:type nums: List[List[int]]
:rtype: List[int]
"""
result, dq, col = [], collections.deque(), 0
for i in range(len(nums) + max(list(map(len, nums)))-1):
new_dq = collections.deque()
if i < len(nums):
dq.appendleft((i, 0))
for r, c in dq:
result.append(nums[r][c])
if c+1 < len(nums[r]):
new_dq.append((r, c+1))
dq = new_dq
return result
[2]:
s = Solution1()
nums = [[1,2,3],[4,5,6],[7,8,9]]
assert s.findDiagonalOrder(nums) == [1,4,2,7,5,3,8,6,9]
nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
assert s.findDiagonalOrder(nums) == [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]
nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]]
assert s.findDiagonalOrder(nums) == [1,4,2,5,3,8,6,9,7,10,11]
nums = [[1,2,3,4,5,6]]
assert s.findDiagonalOrder(nums) == [1,2,3,4,5,6]
[3]:
class Solution2(object):
def findDiagonalOrder(self, nums):
"""
:type nums: List[List[int]]
:rtype: List[int]
"""
result = []
for r, row in enumerate(nums):
for c, num in enumerate(row):
if len(result) <= r+c:
result.append([])
result[r+c].append(num)
return [num for row in result for num in reversed(row)]
[4]:
s = Solution2()
nums = [[1,2,3],[4,5,6],[7,8,9]]
assert s.findDiagonalOrder(nums) == [1,4,2,7,5,3,8,6,9]
nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
assert s.findDiagonalOrder(nums) == [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]
nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]]
assert s.findDiagonalOrder(nums) == [1,4,2,5,3,8,6,9,7,10,11]
nums = [[1,2,3,4,5,6]]
assert s.findDiagonalOrder(nums) == [1,2,3,4,5,6]